#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Sep 05 14:59:47 2021

@author: Kai Lu @ SYSU
"""
# 导入 需要的 library 库
import numpy as np  # 科学计算
import matplotlib.pyplot as plt  # 画图
import matplotlib as mpl
import scipy.signal as sg  # 导入 scipy 的 signal 库 命名为 sg
from scipy.integrate import odeint, solve_bvp, solve_ivp

mpl.rc('lines', linewidth=4)  # , linestyle='-.'
plt.rcParams["font.family"] = "Times New Roman"
plt.rcParams['xtick.labelsize'] = 10
plt.rcParams['ytick.labelsize'] = 10
plt.rcParams['savefig.dpi'] = 300  # 图片像素
plt.rcParams['figure.dpi'] = 300  # 分辨率

# t = np.linspace(-10, 10, num=100, endpoint=True)  # 另一种表达式 t = np.mgrid[0:5:0.01]
# a, b = 2, 0
# gama, omega = 0.1, 1
# C, alpha = a + 1j * b, gama + 1j * omega
# x_t = C * np.exp(alpha * t)
# x_t_re = np.real(x_t)
# x_t_im = np.imag(x_t)
# # y = b * np.exp(a * t) * np.sin(omega * t)
# fig = plt.figure(dpi=300)
# plt.xlabel('Time (s)')
# plt.ylabel('Mag')
# plt.title(r'$Ce^{\alpha t}, C \in \mathbb{R} , \alpha \in \mathbb{C}$: ' + fr'C={C}, $\alpha={alpha}$')
# plt.plot(t, x_t_re, label=r'Real')
# plt.plot(t, x_t_im, label=r'Imag')
# plt.legend(loc='upper left')
# plt.grid(color='k', linestyle='-', linewidth=0.1)
# # plt.grid(axis='x', linewidth=1, linestyle='--', color='0.75')
# # plt.grid(axis='y', linewidth=1, linestyle='--', color='0.75')
# plt.savefig('Signal_exp.png', bbox_inches='tight', pad_inches=0.02, dpi=300)

# 使用方程解
# odeint: Integrate a system of ordinary differential equations
# solve_bvp: Solve a boundary-value problem for a system of ODEs
# solve_ivp: Solve an initial value problem for a system of ODEs

# 一阶微分方程组
def fvdp(t, y):
    '''
    来源：https://www.jianshu.com/p/ab57b600b854?utm_campaign=shakespeare
    要把y看出一个向量，y = [dy0,dy1,dy2,...]分别表示y的n阶导
    对于二阶微分方程，肯定是由0阶和1阶函数组合而成的，所以下面把y看成向量的话，y0表示最初始的函数，也就是我们要求解的函数，y1表示一阶导，对于高阶微分方程也可以以此类推
    '''
    y0, y1 = y
    ft = 10 * np.sin(2 * np.pi * t)
    y2 = -2 * y1 - 77 * y0 + ft
    # y0是需要求解的函数，y1是一阶导
    # 返回的顺序是[一阶导， 二阶导]，这就形成了一阶微分方程组
    return [y1, y2]


y0 = [0, 0]  # 初值[0,0]表示y(0)=0,y'(0)=0
t = np.linspace(0, 5, 100)
y = odeint(fvdp, y0, t, tfirst=True)  # 用 odeint 计算 y(t)
y_ = solve_ivp(fvdp, t_span=(0, 5), y0=y0, t_eval=t)  # 用 solve_ivp 计算 y(t)

# 开始绘图
plt.subplot(211)
y1, = plt.plot(t, y[:, 0], label='y')
y1_, = plt.plot(t, y[:, 1], label='y‘')
plt.legend(handles=[y1, y1_], loc='upper right')
plt.grid(True)

plt.subplot(212)
y2, = plt.plot(y_.t, y_.y[0, :], 'g--', label='y(0)')
y2_, = plt.plot(y_.t, y_.y[1, :], 'r-', label='y(1)')
plt.legend(handles=[y2, y2_], loc='upper right')
plt.grid(True)

plt.show()

# 用已有库的方法解 sg is scipy.signal
sys = sg.lti([1], [1, 2, 77])  # 方程里的系数
ft = 10 * np.sin(2 * np.pi * t)
_, y, _ = sg.lsim(sys, ft, T=t)
# 开始绘图
plt.plot(t, y, label='simple way')
plt.grid(True)
plt.show()

sys = sg.lti([1, 1], [7, 4, 6])  # 方程里的系数 由高次幂到低次幂
st, sy = sg.step2(sys)
it, iy = sg.impulse2(sys)
sy1, = plt.plot(st, sy, label='step')
iy1, = plt.plot(it, iy, label='impluse')
# 开始绘图
plt.legend(handles=[sy1, iy1], loc='upper right')
plt.grid(True)
plt.show()
#
# # sg is scipy.signal
# t1 = np.array([t * 0.1 for t in range(-10, 31)])  # t in [-1, 3]
# f1t = np.array([2 if 0 < t < 10 else 0 for t in range(-10, 31)])
# t2 = np.array([t * 0.1 for t in range(-10, 31)])  # t in [-1,3]
# f2t = np.array([t * 0.1 if 0 < t < 20 else 0 for t in range(-10, 31)])
# yt = sg.convolve(f1t, f2t, 'full') * 0.1  # 计算卷积 calculate convolution
# t3 = np.array([t * 0.1 for t in range(-20, 61)])  # t in [-1+-1, 3+3]
# # 开始绘图
# plt.plot(t3, yt, label='conv')
# plt.grid(True)
# plt.show()
#
# # sg is scipy.signal
# d = np.random.rand(1, 51) - 0.5  # random.rand 出来的是 0到1 的随机数
# k = np.array([k for k in range(0, 51)])
# s = 2 * k * np.power(0.9, k)
# f = s + d[0]
#
# plt.subplot(211)
# plt.stem(k, f, '-', use_line_collection=True)
# plt.grid(True)
#
# M = 5
# a = 1
# b = np.ones(5) / 5
# plt.subplot(212)
# y = sg.filtfilt(b, a, f)  # digital filter forward and backward to a signal
# plt.stem(k, y, ':', use_line_collection=True)
# plt.grid(True)
#
# plt.xlabel('time index k')
# plt.show()
#
# # sg is scipy.signal
# k = np.array([k for k in range(11)])
# a = [1., 3., 2.]
# b = [1.]
# h = sg.lfilter(b, a, k)  # IIR or FIR filter
# plt.stem(k, h, '-', use_line_collection=True)
# plt.grid(True)
# plt.show()
#
# # sg is scipy.signal
# k1 = np.linspace(0, 10, 11)
# x1 = np.sin(k1)
# plt.subplot(221)
# plt.stem(k1, x1, '-', use_line_collection=True)
# plt.grid(True)
# plt.title('x_1(k)=sin(k)')
#
# k2 = np.linspace(0, 15, 16)
# x2 = np.power(0.8, k2)
# plt.subplot(222)
# plt.stem(k2, x2, '-', use_line_collection=True)
# plt.grid(True)
# plt.title('x_2(k) = 0.8^k')
#
# plt.subplot(212)
# y = sg.convolve(x1, x2, 'full')  # 使用 scipy.signal 的卷积函数 convolve
# k3 = np.linspace(0, 25, 26)
# plt.stem(k3, y, '-', use_line_collection=True)
# plt.grid(True)
# plt.title('y(k)')
#
# plt.xlabel('time index k')
# plt.subplots_adjust(top=1, wspace=0.4, hspace=0.5)  # 调整视图
#
# plt.show()